from cogent import *

#
# I assume these filepaths are valid.
#
newickTreePath = "/Users/victor/Documents/THORNTON_LAB/eftu/PHYML-wag/EF_Tu_clustalx_phy1_phyml_tree.tre" # filepath to newick ML tree
partsFilePath = "/Users/victor/Documents/THORNTON_LAB/eftu/MRBAYES/EF-Tu.clustalx.nxs.parts" # filepath to *.parts
trprobsFilePath = "/Users/victor/Documents/THORNTON_LAB/eftu/MRBAYES/EF-Tu.clustalx.nxs.trprobs" # filepath to *.trprobs

#newickTreePath = sys.argv[1]
#partsFilePath = sys.argv[2]
#trprobsFilePath = sys.argv[3]

idsToNames = {} # key = integer, value = string of a taxon name
namesToIds = {} # key = string name of a taxon, value = integer id

f = open(trprobsFilePath, "r")
lines = f.readlines()
foundTranslate = False
for l in lines:
    if l.__contains__("translate"):
        foundTranslate = True
        continue
    if l.__contains__("tree") and foundTranslate == True:
        break
    if foundTranslate == True:
        l = l.replace(",", "")
        l = l.replace(";", "")
        l = l.strip()
        tokens = l.split()
        id = int(tokens[0])
        name = tokens[1]
        idsToNames[id] = name
        namesToIds[name] = id
        print "Translate: " + id.__str__() + " to " + name.__str__()
f.close()

#
# Takes a list of taxa names and returns the hash key for their partition.
# This method is necessary to access elements of the 'partition' data structure,
# where each partition is indexed by names of the descendant taxa.  
#
def taxaToHash(taxa):
    numbers = []
    for t in taxa:
        numbers.append( namesToIds[t] )
    numbers.sort()
    return numbers.__str__()

#
# where partString is a string containing "." and "*"
# returns a list of taxon names
#
def getTaxaInThisPartition(partString):
    returnList = []
    current = 0
    for c in partString:
        current += 1
        if c == "*":
            returnList.append( idsToNames[current])
    return returnList
#
#
#    
partitionProbs = {} # key = partition hash, value = posterior probability of that partition

f = open(partsFilePath, "r")
lines = f.readlines()
foundParts = False
for l in lines:
    if foundParts == False:
        m = re.match("^ID", l)
        if m != None:
            foundParts = True
            continue
    if foundParts == True:
        tokens = l.split()
        part = tokens[1]
        part = re.sub("\d+", "", part)
        prob = float(tokens[2])
        #print part
        taxaInThisPart = getTaxaInThisPartition(part)
        partitionProbs[ taxaToHash(taxaInThisPart) ] = prob
        print part
        print "HASH: " + taxaToHash(taxaInThisPart).__str__() + ": " + prob.__str__()
f.close()

#print "partitionProbs = " + partitionProbs.__str__()

# returns a Newick string
def getNewickTree(tree):
    returnString = ""
    if tree.isTip():
        returnString = tree.Name
        returnString += ":"
        returnString += tree._get_length().__str__()
        return returnString
    else:
        returnString = "("
        for i in range(0, tree.Children.__len__() ):
            c = tree.Children[i]
            returnString += getNewickTree(c)
            if i < tree.Children.__len__() - 1:
                returnString += ","
        returnString += ")"
        if tree.params.__contains__("postprob"):
            returnString += tree.params["postprob"].__str__() + ":" + tree._get_length().__str__()
        else:
            returnString += "unknown:" + tree._get_length().__str__()
    return returnString    

#
# finally, traverse the ML tree add pp support values:
#
root = LoadTree(newickTreePath)
internalNodes = root.nontips()

for n in internalNodes:
    tipNames = []
    tipNodes = n.tips()
    for t in tipNodes:
        tipNames.append( t.Name )
    hashkey = taxaToHash( tipNames )
    if partitionProbs.__contains__(hashkey):
        prob = partitionProbs[ hashkey ]
        print "descendants: " + hashkey + ", " + prob.__str__()
        n.params["postprob"] = prob
    else:
        print "skipping node with hash: " + hashkey
        n.params["postprob"] = 0.0
        
print getNewickTree(root)


